#include "CameraPingPongBuffer.h"
#include <algorithm>

CameraPingPongBuffer::Buffer::Buffer(size_t size) 
{
    data = std::make_unique<uint8_t[]>(size);
}

CameraPingPongBuffer::CameraPingPongBuffer(size_t buffer_size)
    : buffer_size_(buffer_size)
    , buffer_a_(std::make_unique<Buffer>(buffer_size))
    , buffer_b_(std::make_unique<Buffer>(buffer_size))
    , current_write_buffer_(buffer_a_.get())
    , current_read_buffer_(nullptr) {
}

uint8_t* CameraPingPongBuffer::getWriteBuffer() 
{
    std::lock_guard<std::mutex> lock(buffer_mutex_);
    
    // 查找可用的写缓冲区
    Buffer* available_buffer = nullptr;
    
    // 优先使用当前指定的写缓冲区
    if (current_write_buffer_->state.load(std::memory_order_acquire) == BufferState::IDLE) 
    {
        available_buffer = current_write_buffer_;
    } 
    else 
    {
        // 当前写缓冲区不可用，尝试另一个缓冲区
        Buffer* other_buffer = (current_write_buffer_ == buffer_a_.get()) ? 
                              buffer_b_.get() : buffer_a_.get();
        
        if (other_buffer->state.load(std::memory_order_acquire) == BufferState::IDLE) 
        {
            available_buffer = other_buffer;
            current_write_buffer_ = other_buffer;
        }
    }
    
    if (!available_buffer) 
    {
        return nullptr; // 没有可用的写缓冲区
    }
    
    // 原子地将缓冲区状态从IDLE改为WRITING
    BufferState expected = BufferState::IDLE;
    if (!available_buffer->state.compare_exchange_strong(
            expected, BufferState::WRITING, std::memory_order_acq_rel))
    {
        return nullptr; // 状态改变失败，可能被其他线程抢占
    }
    
    return available_buffer->data.get();
}

bool CameraPingPongBuffer::commitWrite(size_t data_size) 
{
    std::lock_guard<std::mutex> lock(buffer_mutex_);
    
    // 查找正在写入状态的缓冲区
    Buffer* writing_buffer = nullptr;
    if (buffer_a_->state.load(std::memory_order_acquire) == BufferState::WRITING) 
    {
        writing_buffer = buffer_a_.get();
    } 
    else if (buffer_b_->state.load(std::memory_order_acquire) == BufferState::WRITING) 
    {
        writing_buffer = buffer_b_.get();
    }
    
    if (!writing_buffer) 
    {
        return false; // 没有找到正在写入的缓冲区
    }
    
    // 限制数据大小，防止越界
    data_size = std::min(data_size, buffer_size_);
    
    // 设置数据大小
    writing_buffer->data_size.store(data_size, std::memory_order_release);
    
    // 原子地将状态从WRITING改为READY
    BufferState expected = BufferState::WRITING;
    if (!writing_buffer->state.compare_exchange_strong(
            expected, BufferState::READY, std::memory_order_acq_rel)) 
    {
        return false; // 状态改变失败
    }
    
    // 为下次写入准备新的写缓冲区
    Buffer* other_buffer = (writing_buffer == buffer_a_.get()) ? 
                          buffer_b_.get() : buffer_a_.get();
    
    // 只有当另一个缓冲区处于IDLE状态时才设为写缓冲区
    if (other_buffer->state.load(std::memory_order_acquire) == BufferState::IDLE) 
    {
        current_write_buffer_ = other_buffer;
    }
    
    return true;
}

const uint8_t* CameraPingPongBuffer::getReadBuffer(size_t& data_size) {
    std::lock_guard<std::mutex> lock(buffer_mutex_);
    
    // 如果当前有正在读取的缓冲区，继续使用它
    if (current_read_buffer_ && 
        current_read_buffer_->state.load(std::memory_order_acquire) == BufferState::READING) 
    {
        data_size = current_read_buffer_->data_size.load(std::memory_order_acquire);
        return current_read_buffer_->data.get();
    }
    
    // 查找READY状态的缓冲区
    Buffer* ready_buffer = nullptr;
    if (buffer_a_->state.load(std::memory_order_acquire) == BufferState::READY) 
    {
        ready_buffer = buffer_a_.get();
    } 
    else if (buffer_b_->state.load(std::memory_order_acquire) == BufferState::READY) 
    {
        ready_buffer = buffer_b_.get();
    }
    
    if (!ready_buffer) 
    {
        return nullptr; // 没有可读数据
    }
    
    // 原子地将状态从READY改为READING
    BufferState expected = BufferState::READY;
    if (!ready_buffer->state.compare_exchange_strong(
            expected, BufferState::READING, std::memory_order_acq_rel)) 
    {
        return nullptr; // 状态改变失败，可能被其他线程抢占
    }
    
    // 清理之前的读缓冲区状态
    if (current_read_buffer_ && current_read_buffer_ != ready_buffer) 
    {
        BufferState old_state = current_read_buffer_->state.load(std::memory_order_acquire);
        if (old_state == BufferState::READING) {
            current_read_buffer_->state.store(BufferState::IDLE, std::memory_order_release);
        }
    }
    
    current_read_buffer_ = ready_buffer;
    data_size = ready_buffer->data_size.load(std::memory_order_acquire);
    return ready_buffer->data.get();
}

void CameraPingPongBuffer::finishRead() 
{
    std::lock_guard<std::mutex> lock(buffer_mutex_);
    
    if (current_read_buffer_ && 
        current_read_buffer_->state.load(std::memory_order_acquire) == BufferState::READING) 
    {
        
        // 原子地将状态从READING改为IDLE
        BufferState expected = BufferState::READING;
        current_read_buffer_->state.compare_exchange_strong(
            expected, BufferState::IDLE, std::memory_order_acq_rel);
        
        current_read_buffer_ = nullptr;
    }
}

bool CameraPingPongBuffer::hasReadableData() const 
{
    std::lock_guard<std::mutex> lock(buffer_mutex_);
    
    return (buffer_a_->state.load(std::memory_order_acquire) == BufferState::READY ||
            buffer_b_->state.load(std::memory_order_acquire) == BufferState::READY ||
            (current_read_buffer_ && 
             current_read_buffer_->state.load(std::memory_order_acquire) == BufferState::READING));
}
